/*
 * Copyright (c) 2025, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.nvidia.spark.rapids

import ai.rapids.cudf.{NvtxColor, NvtxRange}
import java.io.{File, FileOutputStream}
import scala.collection.mutable

sealed case class NvtxId private(name: String, color: NvtxColor, doc: String) {
  def help(): Unit = println(s"$name|$doc")

  def apply[V](block: => V): V = {
    try {
      NvtxRange.pushRange(name, color)
      block
    } finally {
      NvtxRange.popRange()
    }
  }
}

object NvtxRegistry {
  val registeredRanges: mutable.Map[String, NvtxId] = mutable.Map[String, NvtxId]()

  private def register(id: NvtxId): Unit = {
    if (registeredRanges.contains(id.name)) {
      throw new IllegalArgumentException(s"Collision detected for key: ${id.name}")
    } else {
      registeredRanges += (id.name -> id)
    }
  }

  val ACQUIRE_GPU: NvtxId = NvtxId("Acquire GPU", NvtxColor.RED,
      "Time waiting for GPU semaphore to be acquired")

  val RELEASE_GPU: NvtxId = NvtxId("Release GPU", NvtxColor.RED,
    "Releasing the GPU semaphore")

  def init(): Unit = {
    register(ACQUIRE_GPU)
    register(RELEASE_GPU)
  }
}

object NvtxRangeDocs {
  def helpCommon(): Unit = {
    println("---")
    println("layout: page")
    println("title: NVTX Ranges")
    println("nav_order: 5")
    println("parent: Developer Overview")
    println("---")
    println(s"<!-- Generated by NvtxRangeDocs.help. DO NOT EDIT! -->")
    // scalastyle:off line.size.limit
    println("""# RAPIDS Accelerator for Apache Spark Nvtx Range Glossary
              |The following is the list of Nvtx ranges that are used throughout
              |the plugin. To add your own Nvtx range to the code, create an NvtxId
              |entry in NvtxRangeWithDoc.scala and create an `NvtxRangeWithDoc` in the
              |code location that you want to cover, passing in the newly created NvtxId.
              |
              |See [nvtx_profiling.md](https://nvidia.github.io/spark-rapids/docs/dev/nvtx_profiling.html) for more info.
              |
              |""".stripMargin)
    // scalastyle:on line.size.limit
    println("\n## Nvtx Ranges\n")
    println("Name | Description")
    println("-----|-------------")
  }

  def main(args: Array[String]): Unit = {
    NvtxRegistry.init()
    val configs = new FileOutputStream(new File(args(0)))
    Console.withOut(configs) {
      Console.withErr(configs) {
        helpCommon()
        NvtxRegistry.registeredRanges.values.foreach(_.help())
      }
    }
  }
}
